# -*- encoding: utf-8 -*-
"""
========================================
@Time   ：2021/9/7 17:07
@Auther ：shutao
@File   ：QTM_kernel_methods.py
@IDE    ：PyCharm
@Github ：https://github.com/NameLacker
@Gitee  ：https://gitee.com/nameLacker
========================================
"""

import time
import matplotlib
import numpy as np
import paddle
from numpy import pi as PI
from matplotlib import pyplot as plt

from paddle import matmul, transpose
from paddle_quantum.circuit import UAnsatz

import sklearn
from sklearn import svm
from sklearn.datasets import fetch_openml, make_moons, make_circles
from sklearn.model_selection import train_test_split

from IPython.display import clear_output
from tqdm import tqdm


# 初始化一个用来更新进度条的全局变量
N = 1
# 初始化进度条
bar_format_string = '{l_bar}{bar}|[{elapsed}<{remaining}, ' '{rate_fmt}{postfix}]'
pbar = tqdm(total=100, bar_format=bar_format_string)
pbar.close()
clear_output()


def generate_datasets():
    """
    生成训练集和测试集
    Returns:
        返回生成的数据集
    """
    # 生成训练集和测试集
    X_train, y_train = make_circles(10, noise=0.05, factor=0.2)
    X_test, y_test = make_circles(10, noise=0.05, factor=0.2)

    # # 将生成的数据集进行可视化
    # fig, ax = plt.subplots(1, 2, figsize=[10, 4])
    # ax[0].scatter(X_train[:, 0], X_train[:, 1],
    #               marker='o', c=matplotlib.cm.coolwarm(np.array(y_train, dtype=np.float32)))
    # ax[0].set_title('Train')
    # ax[1].set_title('Test')
    # ax[1].scatter(X_test[:, 0], X_test[:, 1], marker='v',
    #               c=matplotlib.cm.coolwarm(np.array(y_test, dtype=np.float32)))
    #
    # print("先来看一下我们的训练集和测试集:")
    # plt.show()
    return X_train, y_train, X_test, y_test


def q_kernel_estimator(x1, x2):
    """
    通过 paddle quantum 模拟 QKE 电路
    Args:
        x1: 数据点1
        x2: 数据点2

    Returns:
        返回测量结果为 0...0 的概率
    """
    x1 = paddle.to_tensor(x1)
    x2 = paddle.to_tensor(x2)

    cir = UAnsatz(2)

    # 添加对应第一个数据的编码电路
    cir.iqp_encoding(x1, pattern=[[0, 1]])

    # 添加对应第二个数据的编码电路的逆电路
    cir.iqp_encoding(x2, pattern=[[0, 1]], invert=True)

    # 用态矢量模式运行电路
    fin_state = cir.run_state_vector()

    # 更新进度条
    global pbar
    global N
    pbar.update(100/N)

    # 返回测量结果为 0...0 的概率
    return (fin_state[0].conj() * fin_state[0]).real().numpy()[0]


def q_kernel_matrix(X1, X2):
    """
    按照规范定义需要传入的 Kernel 矩阵函数
    按照规范，输入应为两个数据列表
    Args:
        X1: 数据列表1
        X2: 数据列表2

    Returns:
        Kernel 矩阵函数
    """
    return np.array([[q_kernel_estimator(x1, x2) for x2 in X2] for x1 in X1])


def visualize_decision_bound(clf, datasets):
    X_train, y_train, X_test, y_test = datasets
    # 创建一个 10x10 的网格
    x_min, x_max = X_train[:, 0].min(), X_train[:, 0].max()
    y_min, y_max = X_train[:, 1].min(), X_train[:, 1].max()
    XX, YY = np.meshgrid(np.linspace(-1.2, 1.2, 10),
                         np.linspace(-1.2, 1.2, 10))

    # 计算分类器该网格上的预测值
    Z = clf.decision_function(np.c_[XX.ravel(), YY.ravel()])
    Z_qke = Z.reshape(XX.shape)

    # 可视化决策函数，以及决策边界
    clear_output()
    plt.contourf(XX, YY, Z_qke, vmin=-1., vmax=1., levels=20,
                 cmap=matplotlib.cm.coolwarm, alpha=1)
    plt.scatter(X_train[:, 0], X_train[:, 1],
                c=matplotlib.cm.coolwarm(np.array(y_train, dtype=np.float32)),
                edgecolor='black')
    plt.scatter(X_test[:, 0], X_test[:, 1], marker='v',
                c=matplotlib.cm.coolwarm(np.array(y_test, dtype=np.float32)),
                edgecolor='black')
    plt.contour(XX, YY, Z_qke, colors=['k', 'k', 'k'], linestyles=['--', '-', '--'],
                levels=[-.2, 0, .2])
    pass


def run_main():
    clear_output()
    X_train, y_train, X_test, y_test = generate_datasets()

    # 作为检查量子核函数是否运行正常，我们可以检查 K(x, x) 是否等于 1
    # 提示：根据内积的定义 <x, x> = 1
    print('检查是否满足 K(x, x) = 1？',
          bool(abs(q_kernel_estimator(np.array([1., 1.]), np.array([1., 1.])) - 1) < 1e-8))

    # 用量子核支持向量机（QKE-SVM）来进行数据分类
    # 创建进度条，并设置所需要的量子核函数计算数量 N
    pbar = tqdm(total=100,
                desc='训练 QKE-SVM 并分类中',
                bar_format=bar_format_string)
    N = len(X_train) ** 2 + len(X_train) ** 2 + len(X_train) * len(X_test)

    # 创建一个具有量子核函数的支持向量机
    svm_qke = svm.SVC(kernel=q_kernel_matrix)

    # 根据训练数据计算支持向量机的决策平面
    svm_qke.fit(X_train, y_train)

    # 计算支持向量机分别对于训练数据和测试数据的分类预测值
    predict_svm_qke_train = svm_qke.predict(X_train)
    predict_svm_qke_test = svm_qke.predict(X_test)

    # 计算准确率
    accuracy_train = np.array(predict_svm_qke_train == y_train, dtype=int).sum() / len(y_train)
    accuracy_test = np.array(predict_svm_qke_test == y_test, dtype=int).sum() / len(y_test)

    # 可视化分类预测结果
    pbar.close()
    clear_output()
    fig, ax = plt.subplots(1, 2, figsize=[10, 4])
    ax[0].scatter(X_train[:, 0], X_train[:, 1], marker='o',
                  c=matplotlib.cm.coolwarm(np.array(predict_svm_qke_train, dtype=np.float32)))
    ax[0].set_title('Prediction on training set, accuracy={:.2f}'.format(accuracy_train))
    ax[1].scatter(X_test[:, 0], X_test[:, 1], marker='v',
                  c=matplotlib.cm.coolwarm(np.array(predict_svm_qke_test, dtype=np.float32)))
    ax[1].set_title('Prediction on testing set, accuracy={:.2f}'.format(accuracy_test))
    print("让我们看一下量子核支持向量机的分类效果如何:")

    # 创建进度条，并设置所需要的量子核函数计算数量 N
    pbar = tqdm(total=100,
                desc='计算 QKE-SVM 的决策函数中',
                bar_format=bar_format_string)
    N = 10 ** 2 * len(X_train)

    # 可视化决策函数平面
    visualize_decision_bound(svm_qke, (X_train, y_train, X_test, y_test))
    pbar.close()


if __name__ == '__main__':
    run_main()
